import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from skimage import io
import torch.nn.functional as F
from train import val_id_list,train_id_list
from models.model import model_Unet_resnet34
from loss import multi_dice_fun

Image_Path = '../data/train_data/resize_image'
Label_Path = '../data/train_data/resize_label'
def show_pred(name,model):
    f, ax = plt.subplots(3, 4, figsize=(8, 8))
    name = name.replace("\n", "")
    image_dir = os.path.join(Image_Path, name)
    label_dir = os.path.join(Label_Path, name.replace("jpg", "npy"))

    image = io.imread(image_dir)
    label = np.load(label_dir)
    image = image[:320,:320,:]
    label = label[:,:320,:320]
    pred = model(((torch.from_numpy(np.transpose(image,(2,0,1))))[np.newaxis,...]).cuda().float())
    pred = (F.sigmoid(pred[0])) > 0.5
    pred = (pred.cpu().numpy()).astype(np.bool)
    print('pred',pred.shape)
    ax[0, 0].imshow(image)
    for i in range(4):
        ax[1,i].imshow(label[i])
    for i in range(4):
        ax[2,i].imshow(pred[i])
    plt.show()


def test(id_list,model):
    num_name = 4*len(id_list)
    dice_all = 0
    for name in id_list:
        name = name.replace("\n", "")
        image_dir = os.path.join(Image_Path, name)
        label_dir = os.path.join(Label_Path, name.replace("jpg", "npy"))

        image = io.imread(image_dir)
        label = (np.load(label_dir)).astype(np.float32)
        image = image[:320, 100:420, :]
        label = label[:, :320, 100:420]
        # print('image',image.shape,'label',label.shape)
        pred = model(((torch.from_numpy(np.transpose(image, (2, 0, 1))))[np.newaxis, ...]).cuda().float())
        pred = (F.sigmoid(pred[0])) > 0.5
        pred = (pred.cpu().numpy()).astype(np.float32)
        dice_1,dice_2,dice_3,dice_4 = multi_dice_fun(pred,label)
        # print('dice_1',dice_1)
        dice_all = dice_all+dice_1+dice_2+dice_3+dice_4
    dice = dice_all/num_name
    print('mean dice is :',dice)



if __name__ == '__main__':
    model_Unet_resnet34.load_state_dict((torch.load('../model_save/Unet_resnet34/148.pth'))['models'])
    # for name in train_id_list:
    #     show_pred(name,models=model_Unet_resnet34)

    test(train_id_list,model_Unet_resnet34)
    test(val_id_list,model_Unet_resnet34)



